"""
Pollinations.AI 图片生成器
支持无限生成，150+ 图像风格，通过 URL 传参
"""

import os
import random
import time
import requests
from typing import Optional, Callable
from urllib.parse import quote
from datetime import datetime


class PollinationsGenerator:
    """Pollinations.AI 图片生成器"""
    
    BASE_URL = "https://image.pollinations.ai/prompt"
    
    # 150+ 图像风格列表
    STYLES = [
        "flux", "flux-realism", "flux-cablyai", "flux-anime", "flux-3d",
        "any-dark", "flux-pro", "turbo", "midjourney", "dall-e",
        "stable-diffusion", "photorealistic", "anime", "cartoon", "3d-render",
        "oil-painting", "watercolor", "sketch", "digital-art", "pixel-art",
        "low-poly", "isometric", "vaporwave", "cyberpunk", "steampunk",
        "fantasy", "sci-fi", "horror", "minimalist", "abstract",
        "impressionism", "surrealism", "pop-art", "art-nouveau", "baroque",
        "renaissance", "gothic", "vintage", "retro", "modern",
        "futuristic", "cinematic", "dramatic", "ethereal", "dreamy",
        "dark", "light", "colorful", "monochrome", "sepia",
        "noir", "pastel", "vibrant", "muted", "high-contrast",
        "soft-focus", "sharp", "bokeh", "macro", "wide-angle",
        "fisheye", "panoramic", "aerial", "underwater", "space",
        "nature", "urban", "industrial", "rustic", "elegant",
        "grunge", "glitch", "neon", "holographic", "metallic",
        "glass", "wood", "stone", "fabric", "paper",
        "chrome", "gold", "silver", "copper", "bronze",
        "crystal", "ice", "fire", "water", "earth",
        "air", "lightning", "smoke", "fog", "rain",
        "snow", "clouds", "sunset", "sunrise", "night",
        "day", "twilight", "autumn", "winter", "spring",
        "summer", "tropical", "desert", "mountain", "ocean",
        "forest", "jungle", "city", "village", "castle",
        "temple", "ruins", "cave", "island", "beach",
        "garden", "park", "street", "alley", "bridge",
        "tower", "dome", "arch", "gate", "door",
        "window", "mirror", "portal", "tunnel", "stairs",
        "corridor", "room", "hall", "chamber", "throne",
        "laboratory", "workshop", "studio", "library", "museum",
        "gallery", "theater", "arena", "stadium", "market",
        "cafe", "restaurant", "bar", "club", "hotel",
        "mansion", "cottage", "cabin", "farmhouse", "lighthouse",
        "windmill", "greenhouse", "observatory", "planetarium", "aquarium",
        "zoo", "circus", "carnival", "fair", "festival",
        "concert", "performance", "exhibition", "installation", "sculpture"
    ]
    
    def __init__(self, save_folder: str = "./pollinations_images", proxies: Optional[dict] = None):
        """
        初始化生成器
        
        Args:
            save_folder: 图片保存根目录路径（内部会创建 completed/ 与 tmp/ 子目录）
            proxies: 可选代理配置，例如 {"http": "socks5h://127.0.0.1:1086", "https": "socks5h://127.0.0.1:1086"}
        """
        self.save_folder = save_folder  # 根目录
        self.proxies = proxies  # 下载图片时使用的代理（默认直连）
        self._ensure_folder_exists()
    
    def _ensure_folder_exists(self):
        """确保保存文件夹及子目录存在"""
        # 根目录
        os.makedirs(self.save_folder, exist_ok=True)
        # 完整图片目录
        os.makedirs(os.path.join(self.save_folder, "completed"), exist_ok=True)
        # 临时文件目录
        os.makedirs(os.path.join(self.save_folder, "tmp"), exist_ok=True)
    
    def generate_image(
        self,
        prompt: str,
        width: int = 576,
        height: int = 1024,
        seed: Optional[int] = None,
        model: str = "flux",
        nologo: bool = True,
        enhance: bool = False,
        progress_callback: Optional[Callable[[str], None]] = None
    ) -> Optional[str]:
        """
        生成单张图片
        
        Args:
            prompt: 图片关键词/描述
            width: 图片宽度
            height: 图片高度
            seed: 随机种子（可选）
            model: 图片模型/风格
            nologo: 是否去除 logo
            enhance: 是否增强提示词
            progress_callback: 进度回调函数
            
        Returns:
            保存的图片路径，失败返回 None
        """
        try:
            # 构建 URL
            encoded_prompt = quote(prompt)
            url = f"{self.BASE_URL}/{encoded_prompt}"
            
            # 添加参数（确保无水印）
            params = {
                "width": width,
                "height": height,
                "model": model,
                "nologo": "true",  # 强制去除 logo
                "nofeed": "true",  # 不显示在公共feed
                "private": "true", # 私密模式
                "enhance": str(enhance).lower()
            }
            
            if seed is not None:
                params["seed"] = seed
            
            if progress_callback:
                progress_callback(f"正在生成图片: {prompt[:50]}...")
            
            # 智能退避重试：对429限流和502/503服务器错误进行重试
            max_backoff_retries = 3
            backoff_retry_count = 0
            request_timeout_seconds = 120
            
            response = None
            while backoff_retry_count <= max_backoff_retries:
                # 内层：超时重试循环
                max_timeout_retries = 3
                for timeout_attempt in range(max_timeout_retries + 1):
                    try:
                        response = requests.get(
                            url,
                            params=params,
                            timeout=request_timeout_seconds,
                            proxies=self.proxies,
                        )
                        break  # 请求成功，跳出超时重试循环
                    except (requests.exceptions.Timeout, requests.exceptions.ConnectionError) as e:
                        if timeout_attempt >= max_timeout_retries:
                            if progress_callback:
                                progress_callback("超时，已跳过")
                            return None
                        if progress_callback:
                            progress_callback(f"超时，{timeout_attempt + 1}/{max_timeout_retries + 1} 重试...")
                        time.sleep(2)
                
                # 如果没有获取到响应（不应该发生），返回 None
                if response is None:
                    return None
                
                # 检查响应状态码
                if response.status_code == 429:
                    # 限流错误 - 使用随机退避避免固定模式
                    if backoff_retry_count < max_backoff_retries:
                        if backoff_retry_count == 0:
                            wait_time = 10  # 第一次：10秒
                        else:
                            wait_time = random.randint(20, 40)  # 后续：20-40秒随机
                        if progress_callback:
                            progress_callback(f"限流，等待{wait_time}秒")
                        time.sleep(wait_time)
                        backoff_retry_count += 1
                        continue  # 继续下一次退避重试
                    else:
                        if progress_callback:
                            progress_callback("限流，已跳过")
                        return None
                
                elif response.status_code == 502 or response.status_code == 503:
                    # 服务器临时错误 - 使用较短的退避时间
                    if backoff_retry_count < max_backoff_retries:
                        wait_time = 5 * (2 ** backoff_retry_count)  # 5秒, 10秒, 20秒
                        if progress_callback:
                            progress_callback(f"{response.status_code}错误，等待{wait_time}秒")
                        time.sleep(wait_time)
                        backoff_retry_count += 1
                        continue  # 继续下一次退避重试
                    else:
                        if progress_callback:
                            progress_callback(f"{response.status_code}错误，已跳过")
                        return None
                
                # 其他错误正常抛出，由外层统一处理成"✗ 生成失败: ..."
                response.raise_for_status()
                
                # 成功获取响应，跳出重试循环
                break
            
            # 生成文件名（使甤13位毫秒时间戳，与样例格式一致）
            timestamp_ms = int(time.time() * 1000)  # 13位毫秒时间戳
            filename = f"{timestamp_ms}.jpeg"
            completed_dir = os.path.join(self.save_folder, "completed")
            tmp_dir = os.path.join(self.save_folder, "tmp")
            os.makedirs(completed_dir, exist_ok=True)
            os.makedirs(tmp_dir, exist_ok=True)
            final_path = os.path.join(completed_dir, filename)
            tmp_path = os.path.join(tmp_dir, filename + ".part")
            
            # 保存图片为 JPEG 格式（先写入临时文件，验证成功后再原子重命名）
            from PIL import Image
            from io import BytesIO
            img_bytes = response.content
            try:
                # 先用 verify 检查图片是否完整、格式正确
                img_for_verify = Image.open(BytesIO(img_bytes))
                img_for_verify.verify()
                # verify 之后需要重新打开一次
                img = Image.open(BytesIO(img_bytes))
                # 转换为 RGB（如果是 RGBA 等模式）
                if img.mode in ("RGBA", "LA", "P"):
                    background = Image.new("RGB", img.size, (255, 255, 255))
                    if img.mode == "P":
                        img = img.convert("RGBA")
                    background.paste(img, mask=img.split()[-1] if img.mode in ("RGBA", "LA") else None)
                    img = background
                elif img.mode != "RGB":
                    img = img.convert("RGB")
                # 先写入临时文件
                img.save(tmp_path, "JPEG", quality=95, optimize=True)
                # 原子替换到最终路径，避免出现半张图片
                os.replace(tmp_path, final_path)
            except ImportError:
                # 如果没有 PIL，只能保证字节写入完整，同样采用临时文件 + 原子重命名
                with open(tmp_path, "wb") as f:
                    f.write(img_bytes)
                os.replace(tmp_path, final_path)
            except Exception:
                # 任意异常都不保留半成品
                try:
                    if os.path.exists(tmp_path):
                        os.remove(tmp_path)
                except Exception:
                    pass
                raise
            
            if progress_callback:
                progress_callback(f"✓ 图片已保存: {filename}")
            
            return final_path
            
        except Exception as e:
            return None
    
    def generate_batch(
        self,
        prompts: list[str],
        width: int = 576,
        height: int = 1024,
        seed: Optional[int] = None,
        model: str = "flux",
        images_per_prompt: int = 1,
        auto_seed: bool = True,
        progress_callback: Optional[Callable[[str], None]] = None,
        should_stop: Optional[Callable[[], bool]] = None
    ) -> list[str]:
        """
        批量生成图片
        
        Args:
            prompts: 关键词列表
            width: 图片宽度
            height: 图片高度
            seed: 基础随机种子
            model: 图片模型/风格
            images_per_prompt: 每个关键词生成的图片数量
            auto_seed: 是否自动生成随机种子
            progress_callback: 进度回调函数
            
        Returns:
            成功保存的图片路径列表
        """
        saved_images = []
        total = len(prompts) * images_per_prompt
        current = 0
        
        if progress_callback:
            progress_callback(f"开始批量生成，共 {total} 张图片")
        
        for prompt in prompts:
            for i in range(images_per_prompt):
                # 检查是否需要停止
                if should_stop and should_stop():
                    if progress_callback:
                        progress_callback(f"⛔ 已停止，已生成 {len(saved_images)}/{total} 张")
                    return saved_images
                
                current += 1
                
                # 确定种子值
                if auto_seed:
                    current_seed = random.randint(1, 999999)
                else:
                    current_seed = seed if seed else random.randint(1, 999999)
                
                if progress_callback:
                    progress_callback(f"进度: {current}/{total} - {prompt[:30]}...")
                
                # 创建错误回调：
                # 1. 把 generate_image 内部的日志转发到主 progress_callback
                # 2. 记录真正的失败信息（带 "✗" 或 "失败"），用于后面追加 [current/total]
                error_msg = None
                def error_callback(msg):
                    nonlocal error_msg
                    if progress_callback:
                        progress_callback(msg)
                    if "✗" in msg or "失败" in msg:
                        error_msg = msg
                
                # 生成图片
                filepath = self.generate_image(
                    prompt=prompt,
                    width=width,
                    height=height,
                    seed=current_seed,
                    model=model,
                    progress_callback=error_callback
                )
                
                if filepath:
                    saved_images.append(filepath)
                    if progress_callback:
                        saved_count = len(saved_images)
                        progress_callback(f"✓ [{saved_count}/{total}] 已保存")
                else:
                    if progress_callback:
                        if error_msg:
                            # 只有在 generate_image 报出明确失败信息时，才追加一行带索引的失败日志
                            progress_callback(error_msg.replace("✗ 生成失败:", f"✗ [{current}/{total}] 失败:"))
                        else:
                            # 比如 429 限流的情况，前面的 error_callback 已经输出了友好的提示，这里就不再重复一行“失败”日志
                            pass
        
        if progress_callback:
            # 使用更清晰的文案，强调是“图片张数”而不是“关键词数”
            progress_callback(
                f"批量生成完成！成功生成 {len(saved_images)}/{total} 张图片"
            )
        
        return saved_images
    
    def generate_random_prompts(self, base_prompt: str, count: int, seed: Optional[int] = None) -> list[str]:
        """
        生成随机关键词变体
        
        Args:
            base_prompt: 基础关键词
            count: 生成数量
            seed: 随机种子（None 则使用真随机）
            
        Returns:
            关键词列表
        """
        if seed is not None:
            random.seed(seed)
        else:
            # 使用真随机，确保每次生成的关键词都不同
            import time
            import os
            # 多重随机源：系统随机 + 时间戳 + 进程号
            true_seed = int.from_bytes(os.urandom(8), 'big') + int(time.time() * 1000000) + os.getpid()
            random.seed(true_seed)
        
        prompts = []
        
        # 随机修饰词（主体风格）
        modifiers = [
            "beautiful", "stunning", "amazing", "gorgeous", "breathtaking",
            "epic", "dramatic", "cinematic", "professional", "high quality",
            "detailed", "intricate", "complex", "elegant", "refined",
            "vibrant", "colorful", "dynamic", "powerful", "majestic",
            "serene", "peaceful", "tranquil", "calm", "soothing",
            "mysterious", "enigmatic", "magical", "fantastical", "surreal",
            "realistic", "photorealistic", "hyper-realistic", "lifelike", "natural",
            "artistic", "creative", "imaginative", "innovative", "unique",
            "modern", "contemporary", "futuristic", "retro", "vintage",
            # 额外扩展一些风格/质感词，增加组合空间
            "dreamlike", "mythical", "celestial", "otherworldly", "otherworldly landscape",
            "luxurious", "ornate", "minimalist", "soft and dreamy", "moody",
            "high contrast", "pastel colored", "vivid", "glossy", "matte",
            "noir style", "cinematic still", "award-winning", "editorial style", "fantasy concept art"
        ]
        
        # 场景词（地点 / 时间）
        scenes = [
            "in nature", "at sunset", "at sunrise", "under moonlight", "in golden hour",
            "in the rain", "in the snow", "in autumn", "in spring", "in summer",
            "in the city", "in the mountains", "by the ocean", "in the forest", "in the desert",
            "in a garden", "in a field", "on a beach", "near a lake", "by a river",
            "in a studio", "on stage", "in a gallery", "in a museum", "in a temple",
            # 新增一些更细的场景
            "on a cliff", "beside a waterfall", "in a misty valley", "on a snowy peak", "in a canyon",
            "in a futuristic city", "in an ancient city", "inside a cathedral", "inside a castle hall", "in a neon-lit street",
            "in a cozy cabin", "inside a glass greenhouse", "in a bamboo forest", "by a lighthouse", "on an island"
        ]

        # 细节 / 镜头 / 画面风格词
        details = [
            "highly detailed", "ultra detailed", "4k resolution", "8k resolution",
            "sharp focus", "soft focus", "with dramatic lighting", "with soft lighting",
            "with volumetric lighting", "backlit", "golden hour lighting", "studio lighting",
            "natural lighting", "cinematic lighting", "rim lighting", "glowing highlights",
            "cinematic composition", "rule of thirds composition", "center composition", "wide shot composition",
            "wide angle lens", "telephoto lens", "macro lens close-up", "fisheye lens effect",
            "long exposure", "bokeh background", "shallow depth of field", "HDR photography", "tilt shift effect",
            "film grain", "analog film look", "digital painting style", "watercolor style", "oil painting style",
            "flat illustration style", "isometric style", "low poly style", "minimalist style", "pixel art style"
        ]
        
        for i in range(count):
            # 随机选择修饰词 + 场景 + 细节
            modifier = random.choice(modifiers)
            scene = random.choice(scenes)
            detail = random.choice(details)
            
            # 组合关键词：主体 + 场景 + 细节
            new_prompt = f"{modifier} {base_prompt} {scene}, {detail}"
            prompts.append(new_prompt)
        
        return prompts
    
    @staticmethod
    def get_available_styles() -> list[str]:
        """获取可用的图像风格列表"""
        return PollinationsGenerator.STYLES.copy()
